﻿/*
	本文件功能：
	IO操作执行体代码
*/

#include "stdafx.h"

/*——————————————————————下面是IRP处理——————————————————————*/

NTSTATUS FbinstRead(IN PDEVICE_OBJECT DeviceObject, IN PIRP IRPMess){
	NTSTATUS status = STATUS_UNSUCCESSFUL;
	//APC临界代码
	FsRtlEnterFileSystem();
	//让程序更强壮一些，SEH的功劳
	__try{
		PIO_STACK_LOCATION irpstack = IoGetCurrentIrpStackLocation(IRPMess);
		PVOID buffer=NULL;
		PMDL Mdl = NULL;
		FBFSVolume *VolumeInfo = NULL;
		if (DeviceObject == pDriverExtension->ControlDevice)
		{
			EasyCompleteIRP(IRPMess,STATUS_SUCCESS,0);
			return STATUS_SUCCESS;
		}else{
			UD_Ext_Flag *Ext = (UD_Ext_Flag*)DeviceObject->DeviceExtension;
			VolumeInfo = Ext->Volume;
		}
		//懒得处理DPC异步读了
		if (irpstack->MinorFunction==IRP_MN_DPC)
		{
			EasyCompleteIRP(IRPMess,STATUS_NOT_IMPLEMENTED,0);
			return STATUS_NOT_IMPLEMENTED;
		}
		PAGED_CODE();
		//IRP中传达过来的文件名类似\Device\UDVolume1\aa.dat
		if (irpstack->MinorFunction==IRP_MN_MDL){
			buffer = ExAllocatePool(NonPagedPool,irpstack->Parameters.Read.Length);
			if (buffer == NULL)
			{
				EasyCompleteIRP(IRPMess, STATUS_INSUFFICIENT_RESOURCES,0);
				return STATUS_INSUFFICIENT_RESOURCES;
			}
			Mdl = IoAllocateMdl(buffer,irpstack->Parameters.Read.Length,FALSE,FALSE,IRPMess);
			MmBuildMdlForNonPagedPool(Mdl);
			IRPMess->MdlAddress = Mdl;
			//FbinstFastIoRead(irpstack->FileObject,&irpstack->Read.ByteOffset,irpstack->Read.Length,TRUE,irpstack->Read.Key,buffer,&IRPMess->IoStatus,DeviceObject);
		}else{
			buffer = MmGetSystemAddressForMdl(IRPMess->MdlAddress);
		}
		if (irpstack->MinorFunction==IRP_MN_COMPLETE_MDL){
			buffer = MmGetSystemAddressForMdl(IRPMess->MdlAddress);
			IoFreeMdl(IRPMess->MdlAddress);
			ExFreePool(buffer);
			EasyCompleteIRP(IRPMess,STATUS_SUCCESS,0);
			return STATUS_SUCCESS;
		}
		UINT i = 0;
		
		wchar_t *FileName = wcsncpy((wchar_t*)ExAllocatePool(NonPagedPool,sizeof(wchar_t) * (irpstack->FileObject->FileName.Length+1)),irpstack->FileObject->FileName.Buffer+18,irpstack->FileObject->FileName.Length);
		KeWaitForMutexObject(&pDriverExtension->udListLock,Executive,KernelMode,FALSE,NULL);
		FileRecord * FileHan = NULL;
		for (PLIST_ENTRY p = VolumeInfo->FileListHead.Flink; p != &VolumeInfo->FileListHead;p = p->Flink)
		{
			FileHan = CONTAINING_RECORD(p,FileRecord,list_item);
			if (!wcsncmp(FileHan->FileName[1].Buffer,FileName + wcslen(L"\\Device\\UDVolume1\\"),FileHan->FileName[1].Length)){
				break;
			}else{
				FileHan = NULL;
			}
		}
		KeReleaseMutex(&pDriverExtension->udListLock,FALSE);
		if (FileHan == NULL)
		{
			EasyCompleteIRP(IRPMess,STATUS_OBJECT_NAME_NOT_FOUND,0);
			return STATUS_OBJECT_NAME_NOT_FOUND;
		}
		IO_STATUS_BLOCK io_stk;
		status = FBFSReadFile(*FileHan,buffer,MmGetMdlByteCount(IRPMess->MdlAddress),irpstack->Parameters.Read.Length,irpstack->Parameters.Read.ByteOffset,io_stk);
		if (!NT_SUCCESS(status))
		{
			ExFreePool((void*)FileName);
			IRPMess->IoStatus.Status = status;
			IRPMess->IoStatus.Information = 0;
			IoCompleteRequest(IRPMess,IO_NO_INCREMENT);
			FsRtlExitFileSystem();
			return status;
		}
		ExFreePool((void*)FileName);
		EasyCompleteIRP(IRPMess,status,io_stk.Information);
	}__except(EXCEPTION_EXECUTE_HANDLER){
		EasyCompleteIRP(IRPMess,status,0);
	};
	FsRtlExitFileSystem();
	return status;
}

NTSTATUS FBFSReadFile(FileRecord & FileHandle,void *buffer,size_t bufsize,size_t readsize,LARGE_INTEGER offest,IO_STATUS_BLOCK &IoStatus)
{
	NTSTATUS status;
	IoStatus.Information = 0;
	__try{
		if (bufsize < readsize)
		{
			status = STATUS_BUFFER_TOO_SMALL;
			return status;
		}
		LARGE_INTEGER offest;
		offest.QuadPart = FileHandle.offest.QuadPart;
		if (offest.QuadPart < FileHandle.OnDevice->primary_size)
		{
			PVOID buffer510 = ExAllocatePool(NonPagedPool,510);
			//关于结尾没有扇区对其的情况做特殊处理
			if ((FileHandle.FileLength % 512) != 0)
			{
				unsigned __int64 i = 0;
				for (;i<(readsize / 512);i++)
				{
					PIRP oprirp;
					IO_STATUS_BLOCK io_stk;
					KEVENT Event;
					KeInitializeEvent(&Event,SynchronizationEvent,FALSE);
					oprirp = IoBuildAsynchronousFsdRequest(IRP_MJ_READ,FileHandle.OnDevice->VolumeDevice,buffer,510,&offest,&io_stk);
					IoSetCompletionRoutine(oprirp,&Completion,&Event,TRUE,TRUE,TRUE);
					status = IoCallDriver(FileHandle.OnVolume->device,oprirp);
					if (!NT_SUCCESS(status))
						KeWaitForSingleObject(&Event,Executive,KernelMode,FALSE,NULL);
					if (!NT_SUCCESS(status))
					{//引发一个异常，跳到except块去处理
						*except = 0;
					}
					offest.QuadPart += 512;
					memcat(buffer,buffer510,i,510);
					IoStatus.Information += 510;
				}
				//status = ZwReadFile(FileHandle.OnVolume->CurrentDiskHandle,NULL,NULL,NULL,&IoStatus,buffer510,FileHandle.FileLength % 512,&offest,NULL);
				if (!NT_SUCCESS(status))
					*except = 0;
				memcat(buffer,buffer510,i,FileHandle.FileLength % 512);
				IoStatus.Information += FileHandle.FileLength % 512;
			}else{
				for (unsigned __int64 i = 0;i<readsize;i++)
				{
					//status = ZwReadFile(FileHandle.OnVolume->CurrentDiskHandle,NULL,NULL,NULL,&IoStatus,buffer,510,&offest,NULL);
					PIRP oprirp;
					KEVENT Event;
					IO_STATUS_BLOCK io_stk;
					KeInitializeEvent(&Event,SynchronizationEvent,FALSE);
					oprirp = IoBuildAsynchronousFsdRequest(IRP_MJ_READ,FileHandle.OnDevice->VolumeDevice,buffer,510,&offest,&io_stk);
					IoSetCompletionRoutine(oprirp,&Completion,&Event,TRUE,TRUE,TRUE);
					IoCallDriver(FileHandle.OnVolume->device,oprirp);
					KeWaitForSingleObject(&Event,Executive,KernelMode,FALSE,NULL);
					IoFreeIrp(oprirp);
					if (!NT_SUCCESS(status))
						*except = 0;
					offest.QuadPart += 512;
					memcat(buffer,buffer510,i,510);
				}
				ExFreePool(buffer);
			}
		}else{
			//拓展分区的则按一揽子处理
			status = ZwReadFile(FileHandle.OnVolume->CurrentDiskHandle,NULL,NULL,NULL,&IoStatus,buffer,readsize,&offest,NULL);
			if (!NT_SUCCESS(status))
				*except = 0;
		}
		readsize = (ULONG)IoStatus.Information;
		return status;
	}__except(EXCEPTION_EXECUTE_HANDLER){
		return status;
	}
	return status;
}


NTSTATUS FbinstWrite(IN PDEVICE_OBJECT DeviceObject, IN PIRP IRPMess){
	//懒得处理写了，也为了安全，直接返回拒绝访问了。
	NTSTATUS status = STATUS_ACCESS_DENIED;
	IRPMess->IoStatus.Status = status;
	IRPMess->IoStatus.Information = 0;
	IoCompleteRequest(IRPMess,IO_NO_INCREMENT);
	return status;
}

NTSTATUS FBFSSetFileInfo(IN PDEVICE_OBJECT DeviceObject, IN PIRP IRPMess){
	NTSTATUS status = STATUS_ACCESS_DENIED;
	IRPMess->IoStatus.Status = status;
	IRPMess->IoStatus.Information = 0;
	IoCompleteRequest(IRPMess,IO_NO_INCREMENT);
	return status;
}

NTSTATUS FBFSQueryFileInfo(IN PDEVICE_OBJECT DeviceObject, IN PIRP IRPMess){
	PIO_STACK_LOCATION irpstack = IoGetCurrentIrpStackLocation(IRPMess);
	__try{
		FBFSVolume *VolumeInfo = NULL;
		FileRecord *FileHandle = NULL;
		UNICODE_STRING FileName;
		if (DeviceObject == pDriverExtension->ControlDevice)
		{
			EasyCompleteIRP(IRPMess,STATUS_SUCCESS,0);
			return STATUS_SUCCESS;
		}else{
			UD_Ext_Flag *Ext = (UD_Ext_Flag*)DeviceObject->DeviceExtension;
			VolumeInfo = Ext->Volume;
		}
		UNICODE_STRING UFileName = RTL_CONSTANT_STRING(irpstack->FileObject->FileName.Buffer+wcslen(L"\\Device\\UDVolume1\\"));
		for (PLIST_ENTRY p = VolumeInfo->FileListHead.Flink;p != &VolumeInfo->FileListHead;p = p->Flink)
		{
			FileHandle = CONTAINING_RECORD(p,FileRecord,list_item);
			if (!RtlEqualUnicodeString(&FileHandle->FileName[1],&FileName,FALSE))
			{
				break;
			}else{
				FileHandle = NULL;
			}
		}
		switch (irpstack->Parameters.QueryFile.FileInformationClass)
		{
			case FileAllInformation:
				{
					FILE_ALL_INFORMATION * buffer = (FILE_ALL_INFORMATION*)IRPMess->UserBuffer;
					buffer->AccessInformation.AccessFlags = GENERIC_READ|STANDARD_RIGHTS_READ;
					buffer->AlignmentInformation.AlignmentRequirement = FILE_BYTE_ALIGNMENT;
					buffer->BasicInformation.ChangeTime.QuadPart = FileHandle->ModTime.KeModTime.QuadPart;
					buffer->BasicInformation.CreationTime.QuadPart = FileHandle->ModTime.KeModTime.QuadPart;
					buffer->BasicInformation.FileAttributes = FILE_ATTRIBUTE_READONLY;
					if (FileHandle->DirStruct.IsDir == true)
					{
						buffer->BasicInformation.FileAttributes |= 0;
					}
					break;
				}
			case FileAttributeTagInformation:
			{
				FILE_ATTRIBUTE_TAG_INFORMATION *retdata = (FILE_ATTRIBUTE_TAG_INFORMATION *)IRPMess->AssociatedIrp.SystemBuffer;
				retdata->FileAttributes = FILE_ATTRIBUTE_READONLY;

			}
		}
	}__except(EXCEPTION_EXECUTE_HANDLER){
		NTSTATUS status = STATUS_UNSUCCESSFUL;
		IRPMess->IoStatus.Status = status;
		IRPMess->IoStatus.Information = 0;
		IoCompleteRequest(IRPMess,IO_NO_INCREMENT);
		return status;
	}
}


NTSTATUS Completion(PDEVICE_OBJECT DeviceObject,PIRP IRPMess,PVOID context){
	NTSTATUS status = STATUS_CONTINUE_COMPLETION;
	__try{
		PKEVENT Event = (PKEVENT)context;
		ASSERT(context);
		KeSetEvent(Event,IO_NO_INCREMENT,FALSE);
	}__except(EXCEPTION_EXECUTE_HANDLER){
		status = STATUS_INVALID_PARAMETER;
	}
	IoFreeMdl(IRPMess->MdlAddress);
	IoFreeIrp(IRPMess);
	return status;
}

NTSTATUS FBFSPnP(IN PDEVICE_OBJECT DeviceObject, IN PIRP IRPMess){
	PIO_STACK_LOCATION irpstack = IoGetCurrentIrpStackLocation(IRPMess);
	NTSTATUS status = STATUS_UNSUCCESSFUL;
	__try{
		switch (irpstack->MinorFunction)
		{
		case IRP_MN_REMOVE_DEVICE:
			UD_Ext_Flag *Flag = (UD_Ext_Flag*)irpstack->DeviceObject->DeviceExtension;
			if (Flag->Signature==Volume_Sign)
			{
				KeWaitForSingleObject(&pDriverExtension->udListLock,Executive,KernelMode,FALSE,NULL);
				DiskListEntry *elem = CONTAINING_RECORD(SearchVolume(Flag->Volume),DiskListEntry,udlist_item);
				ZwClose(elem->CurrentDiskHandle);
				UD_Ext_Flag *LowerDevice = (UD_Ext_Flag*)DeviceObject->DeviceExtension;
				IO_STATUS_BLOCK io_stk;
				KEVENT SyncEvent;
				KeInitializeEvent(&SyncEvent,SynchronizationEvent,FALSE);
				PIRP ControlIRP = IoBuildDeviceIoControlRequest(IOCTL_DELETE_DISK,LowerControl,LowerDevice->LowerDevice,4,NULL,0,FALSE,&SyncEvent,&io_stk);
				IoCallDriver(LowerControl,ControlIRP);

			} 
			else
			{
				EasyCompleteIRP(IRPMess,status,0);
			}
			break;
		}
	}__except(EXCEPTION_EXECUTE_HANDLER){
		EasyCompleteIRP(IRPMess,status,0);
	}
	return status;
}